import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
import { OllamaEmbeddings } from "@langchain/ollama";
import { Document } from "@langchain/core/documents";

// 1. 初始化向量存储（复用你的数据库配置，全局单例避免重复连接）
let vectorStore: PGVectorStore | null = null;

async function initVectorStore() {
  if (!vectorStore) {
    try {
      console.log("🔄 正在初始化向量数据库连接...");
      vectorStore = await PGVectorStore.initialize(
        new OllamaEmbeddings({
          baseUrl: 'http://192.168.31.159:11434',
          model: 'quentinz/bge-large-zh-v1.5:latest',
        }),
        {
          tableName: 'embedding_zqw_1024', // 数据库中的表名
          columns: {
            idColumnName: 'embedding_id', // 主键列
            vectorColumnName: 'embedding', // 向量列
            contentColumnName: 'text', // 文本列
            metadataColumnName: 'metadata', // 元数据列
          },
          postgresConnectionOptions: {
            host: '118.178.106.254',
            port: 15432,
            user: 'postgres',
            password: '6a29c7071ae1f1b9',
            database: 'postgres',
          },
        }
      );
      console.log("✅ 向量数据库连接初始化完成");
    } catch (error: unknown) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      console.error("❌ 向量数据库连接失败:", error);
      throw new Error(`向量数据库连接失败: ${errorMessage}`);
    }
  }
  return vectorStore;
}

// 定义检索结果类型（与 state.ts 中的 retrievedDocs 字段匹配）
export type RetrievedDoc = {
  content: string; // 文档内容
  score: number; // 相似度分数（0-1，越高越相关）
  source: string; // 文档来源（从 metadata 提取）
};

// 3. 实现多种检索器（适配不同策略）
/**
 * 基础检索器：通用逻辑，转换文档格式为 RetrievedDoc
 */
async function baseRetrieve(
  question: string,
  retriever: ReturnType<PGVectorStore["asRetriever"]>
): Promise<RetrievedDoc[]> {
  // 调用检索器获取原始文档（LangChain 的 Document 类型）
  const docs: Document[] = await retriever.invoke(question);
  // 转换为 RAG 状态需要的格式
  return docs.map((doc) => ({
    content: doc.pageContent,
    score: (doc.metadata._score as number) || 0.5, // 向量数据库返回的相似度分数，如果没有则默认0.5
    source: doc.metadata.source || "unknown", // 从 metadata 取来源（需确保数据库中存储了 source）
  }));
}

/**
 * 1. TopK 检索器：按相似度取前 K 篇文档（最常用）
 * @param question 用户问题
 * @param k 取前 k 篇（默认 3）
 */
// 最基础的检索器，按相似度取前 K 篇
export async function topKRetriever(
  question: string,
  k: number = 3
): Promise<RetrievedDoc[]> {
  const vectorStore = await initVectorStore();
  // 创建 TopK 检索器（LangChain 内置，按相似度排序取前 k）
  const retriever = vectorStore.asRetriever(k);
  const result = await baseRetrieve(question, retriever);
  console.log(`🔍 TopK 检索完成（k=${k}），返回 ${result.length} 篇文档`);
  return result;
}

/**
 * 2. 混合检索器：结合向量检索和关键词过滤（例如按时间/来源筛选）
 * @param question 用户问题
 * @param filter 关键词过滤条件（如 { source: "老周的修鞋铺" }）
 */

// 支持按元数据过滤（例如只检索特定来源的文档），适合需要限定范围的场景；
export async function hybridRetriever(
  question: string,
  filter: Record<string, unknown> = {}
): Promise<RetrievedDoc[]> {
  const vectorStore = await initVectorStore();
  // 创建带过滤条件的检索器（向量检索 + 元数据过滤）
  const retriever = vectorStore.asRetriever(3, {
    filter, // 例如：只检索来源为 "老周的修鞋铺" 的文档
  });
  const result = await baseRetrieve(question, retriever);
  console.log(`🔍 混合检索完成（过滤条件：${JSON.stringify(filter)}），返回 ${result.length} 篇文档`);
  return result;
}

/**
 * 3. 时间感知检索器：优先返回最新文档（需数据库中 metadata 包含 timestamp 字段）
 * @param question 用户问题
 * @param days 只检索最近 N 天的文档
 */
// 优先最新文档（需你的向量表 metadata 中包含 timestamp 字段，可根据实际情况调整）。
export async function timeAwareRetriever(
  question: string,
  days: number = 30
): Promise<RetrievedDoc[]> {
  const vectorStore = await initVectorStore();
  // 计算 N 天前的日期，格式化为 "YYYY-MM-DD" 字符串
  const cutoffDate = new Date();
  cutoffDate.setDate(cutoffDate.getDate() - days);
  const cutoffDateStr = cutoffDate.toISOString().split('T')[0]; // 示例："2025-11-01"

  // 注意：PGVectorStore的filter只支持简单的等值匹配，不支持范围查询
  // 如果需要时间过滤，建议在检索后手动过滤，或者使用数据库视图
  // 这里先检索更多文档，然后在结果中按时间排序
  const retriever = vectorStore.asRetriever(10); // 多检索一些，后续可以按时间排序
  const result = await baseRetrieve(question, retriever);
  
  // 如果有Date字段，可以在这里进行时间过滤和排序
  // 注意：这需要baseRetrieve返回的文档包含Date信息
  // 目前简化处理，直接返回前3个结果
  // cutoffDateStr 可用于后续的时间过滤逻辑（如果需要）
  const filtered = result.slice(0, 3);
  
  console.log(`🔍 时间感知检索完成（最近 ${days} 天，截止日期：${cutoffDateStr}），返回 ${filtered.length} 篇文档`);
  return filtered;
}